from PyQt6.QtCore import QThread, pyqtSignal
from ultralytics import YOLO
import traceback
import os

class DetectionWorker(QThread):
    finished = pyqtSignal(object)
    progress = pyqtSignal(int)

    def __init__(self, model_path, image_path):
        super().__init__()
        self.model_path = model_path
        self.image_path = image_path

    def run(self):
        try:
            print(f"开始检测: model_path={self.model_path}, image_path={self.image_path}")
            
            # 检查文件是否存在
            if not os.path.exists(self.model_path):
                raise FileNotFoundError(f"模型文件不存在: {self.model_path}")
            
            if not os.path.exists(self.image_path):
                raise FileNotFoundError(f"图像文件不存在: {self.image_path}")
            
            self.progress.emit(10)
            model = YOLO(self.model_path)
            self.progress.emit(30)
            results = model.predict(source=self.image_path, save=False, show=False)
            self.progress.emit(80)

            result = results[0]
            boxes = result.boxes
            num_detections = len(boxes) if boxes is not None else 0

            annotated_img = result.plot(conf=False, labels=False)

            detection_result = {
                "num_detections": num_detections,
                "annotated_image": annotated_img,
                "results": result
            }

            self.progress.emit(100)
            self.finished.emit(detection_result)

        except Exception as e:
            # 提供更详细的错误信息
            error_details = f"{str(e)}\n{traceback.format_exc()}"
            print(f"检测过程中发生错误: {error_details}")  # 添加调试信息
            self.finished.emit({"error": error_details})